# -*- coding: utf-8 -*-
# @Time    : 2021/4/30 16:36
# @Author  : huangwei
# @File    : bxdmethod.py
# @Software: PyCharm
import copy

import cv2
import numpy as np
from PIL import Image
from pyzbar import pyzbar
from word import TextDetector, TextRecognizer


# 识别二维码信息
def code_info(img):
    # 可以同时识别多个二维码
    codes1 = pyzbar.decode(img, symbols=[pyzbar.ZBarSymbol.QRCODE])

    # 进行灰度化和增加对比度，使识别效果更好
    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    kernel = np.array([[0, -1, 0], [-1, 5, -1], [0, -1, 0]], np.float32)
    img = cv2.filter2D(img, -1, kernel)
    codes2 = pyzbar.decode(img, symbols=[pyzbar.ZBarSymbol.QRCODE])

    codes = codes1 + codes2
    info_list = []
    for code in codes:
        # 提取二维码数据为字节对象
        codedata = code.data.decode('utf-8')
        if len(info_list) == 0:
            info_list.append(codedata)
        else:
            for info in info_list:
                if info != codedata:
                    info_list.append(codedata)

    return info_list


# 叉乘
def get_cross(x1, y1, x2, y2, x, y):
    a = (x2 - x1, y2 - y1)
    b = (x - x1, y - y1)

    cross = a[0] * b[1] - a[1] * b[0]

    return cross


# 点是否在box中
def in_box(point, box):
    # 判断点是否在box中
    x, y = point
    x1, y1, x2, y2, x3, y3, x4, y4 = box

    # 使用叉乘法判断
    line12 = get_cross(x1, y1, x2, y2, x, y)
    line34 = get_cross(x4, y4, x3, y3, x, y)

    line14 = get_cross(x1, y1, x4, y4, x, y)
    line23 = get_cross(x2, y2, x3, y3, x, y)

    if line12 * line34 < 0 and line14 * line23 < 0:
        return True

    return False


# 旋转剪切图片
def get_rotate_crop_image(img, points):
    img_crop_width = int(
        max(
            np.linalg.norm(points[0] - points[1]),
            np.linalg.norm(points[2] - points[3])))
    img_crop_height = int(
        max(
            np.linalg.norm(points[0] - points[3]),
            np.linalg.norm(points[1] - points[2])))
    pts_std = np.float32([[0, 0], [img_crop_width, 0],
                          [img_crop_width, img_crop_height],
                          [0, img_crop_height]])
    M = cv2.getPerspectiveTransform(points, pts_std)
    dst_img = cv2.warpPerspective(
        img,
        M, (img_crop_width, img_crop_height),
        borderMode=cv2.BORDER_REPLICATE,
        flags=cv2.INTER_CUBIC)
    dst_img_height, dst_img_width = dst_img.shape[0:2]
    if dst_img_height * 1.0 / dst_img_width >= 1.5:
        dst_img = np.rot90(dst_img)
    return dst_img


# 对box进行从上到下，从左到右的排序
def sorted_boxes(dt_boxes):
    """
    Sort text boxes in order from top to bottom, left to right
    args:
        dt_boxes(array):detected text boxes with shape [4, 2]
    return:
        sorted boxes(array) with shape [4, 2]
    """
    num_boxes = dt_boxes.shape[0]
    sorted_boxes = sorted(dt_boxes, key=lambda x: (x[0][1], x[0][0]))
    _boxes = list(sorted_boxes)

    for i in range(num_boxes - 1):
        if abs(_boxes[i + 1][0][1] - _boxes[i][0][1]) < 10 and \
                (_boxes[i + 1][0][0] < _boxes[i][0][0]):
            tmp = _boxes[i]
            _boxes[i] = _boxes[i + 1]
            _boxes[i + 1] = tmp
    return _boxes


# 识别图片中的文字和文字框
class TextSystem(object):
    def __init__(self, args):
        self.text_detector = TextDetector(args)
        self.text_recognizer = TextRecognizer(args)

    def __call__(self, img):
        ori_im = img.copy()
        # 识别出有文字的框
        dt_boxes = self.text_detector(img)
        print("识别出{}个有文字的框".format(len(dt_boxes)))

        if dt_boxes is None:
            return None, None

        # 对box进行从上到下从左到右的排序
        dt_boxes = sorted_boxes(dt_boxes)

        img_crop_list = []
        # 将识别出的框裁剪出来，再进行文字识别
        for bno in range(len(dt_boxes)):
            tmp_box = copy.deepcopy(dt_boxes[bno])
            img_crop = get_rotate_crop_image(ori_im, tmp_box)
            img_crop_list.append(img_crop)

        # 识别框中的文字
        rec_res = self.text_recognizer(img_crop_list)

        # 返回识别率较高的框和结果
        total_result = []

        # 返回框和识别的结果，一一对应
        for box, rec_reuslt in zip(dt_boxes, rec_res):
            line = {}
            text, score = rec_reuslt
            if score >= 0.5:
                line["box"] = box
                line["rec"] = rec_reuslt
                total_result.append(line)
        return total_result


# 图片过大则缩放
def normal_size(img):
    height, width = img.shape[0:2]
    if height > 2000 or width > 2000:
        if height > width:
            new_height = 2000
            new_width = int(new_height / height * width)
        else:
            new_width = 2000
            new_height = int(new_width / width * height)
        img = cv2.resize(img, (new_width, new_height))
    return img
